#include <iostream>
#include <string>
#include <fstream>
#include <vector>
#include <sstream>
#include <map>
#include <utility>
#include <algorithm>
#include <set>
#include <exception>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include "string_manager.h"
#include "kaldi_scp_process.h"
#include "pevents.h"
#include "pevents_suppliment.h"



#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/lexical_cast.hpp>
using namespace boost::algorithm;
#include <boost/regex.hpp>
#ifdef _MSC_VER
#include <windows.h>
#include <direct.h>
#else
#include <unistd.h>
#endif





using namespace std;


typedef std::map<string, string> ScpMap;
//define some global vars
typedef std::set<string> StringSet;
typedef std::vector<string> StringArray;
typedef std::unordered_map<string, string> FastMap;

//global vars for multi thread processing
static FastMap utt2spk_map, utt2feats_map, utt2targets_map, utt2spk_content;
static StringArray utt2spk_array, utt2spk_keys_array;
static StringArray sub_tasks; 

//for data safeguard
static neosmart::neosmart_event_t h_event_start;
static neosmart::neosmart_event_t h_event_stop;
static neosmart::neosmart_event_t h_event_ack;

static string sub_task_file_name = "";
static int total_task_num = 0;
using namespace neosmart;


bool MakeSureDirsExists(const char* char_dir){
	if (CheckFile(char_dir)){
		return true;
	}
	boost::filesystem::path path(char_dir);
	boost::filesystem::create_directories(path);
	return true;
}


typedef struct _SubTaskInfo
{
	
	int sub_task_index_;
	string sub_task_root_;
	bool is_ready_;
	_SubTaskInfo(){
		sub_task_index_ = -1;
		sub_task_root_ = "";
		is_ready_ = false;
	}
	_SubTaskInfo(string dir, int index, bool ready=false){
		sub_task_index_ = index;
		sub_task_root_ = dir;
		is_ready_ = ready;
	};
	~_SubTaskInfo(){
		sub_task_root_ = "";
		sub_task_index_ = -1;
		is_ready_ = false;
	}
	void Reset(){
		sub_task_root_ = "";
		sub_task_index_ = -1;
		is_ready_ = false;
	}
	void SetState(bool state = false){
		is_ready_ = state;
	}

	bool IsReady(){
		return is_ready_;
	}
}SubTaskInfo;

static SubTaskInfo task_info;


//////////////////////////////////////////////////////////////////////////
int producer_consumer_init(){
	h_event_start = neosmart::CreateEvent(true, false);
	h_event_stop = neosmart::CreateEvent(true,false);
	h_event_ack = neosmart::CreateEvent(true, false);
	return 0;
}


int producer_consumer_reset(){
	if (h_event_start || h_event_stop){
		neosmart::ResetEvent(h_event_start);
		neosmart::ResetEvent(h_event_stop);
		neosmart::ResetEvent(h_event_ack);
	}
	return 0;
}


int producer_consumer_stop(){
	if (h_event_start || h_event_stop){
		neosmart::DestroyEvent(h_event_start);
		neosmart::DestroyEvent(h_event_stop);
		neosmart::DestroyEvent(h_event_ack);
	}
	h_event_start = nullptr;
	h_event_stop = nullptr;
	h_event_ack = nullptr;
	return 0;
}

//////////////////////////////////////////////////////////////////////////


int write_splitted_file(const char* splitted_file_name, const StringArray* sub_task_entries,  FastMap* all_map, int& good_num){
	if (!splitted_file_name || !sub_task_entries || ! all_map || sub_task_entries->size() == 0 ||
		all_map->size() == 0){
		fprintf(stderr, "Warning:%s is not splitted.", splitted_file_name);
		return -1;
	}
	int succ_num = 0;
	string file_name_open = splitted_file_name;
	std::ofstream scp_os(file_name_open.c_str(), ios::out);
	if (scp_os.good()){
		StringArray::const_iterator it = sub_task_entries->begin(), end = sub_task_entries->end();
		for (; it != end; it++){
			string curr_key = *it;
			string curr_content = (*all_map)[curr_key];
			string entry = curr_key + " " + curr_content;
			if (file_name_open.substr(file_name_open.size()-11) == "targets.scp"){
				string fixed_curr_content = FixPdf(curr_content.c_str());
				entry.clear();
				entry  = curr_key + " " + fixed_curr_content;
			}
			if (curr_key != "" && curr_content != ""){
				scp_os << entry << std::endl;
				succ_num++;
			}
		}
	}
	scp_os.close();
	good_num = succ_num;
	return 0;
}



int WriteScpFile(){
	neosmart::neosmart_event_t h_events[2];
	h_events[0] = h_event_start;
	h_events[1] = h_event_stop;
	int succ_num = 0;
	while (true)
	{
		neosmart::WaitForMultipleEvents(h_events, sizeof(h_events) / sizeof(neosmart::neosmart_event_t), false, -1);
		int succ_num_one_time=0, succ_num_one_time1 = 0, succ_num_one_time2 = 0, succ_num_one_time3 = 0;
		StdThread data_protect;
		if (neosmart::CheckEvent(h_event_start)){
			//do something,when data is ready, always start event
			if (task_info.IsReady()){
				string sub_dir, curr_dir = task_info.sub_task_root_;
				stringstream tmp_ss;
				tmp_ss << task_info.sub_task_index_;
				tmp_ss >> sub_dir;
				string base_dir = curr_dir + "\/" + sub_dir;
				MakeSureDirsExists(base_dir.c_str());
				string sub_feats = base_dir + "\/" + "feats.scp",
					sub_targets = base_dir + "\/" + "targets.scp",
					sub_utt2spk = base_dir + "\/" + "utt2spk";
				data_protect.Enter();
				write_splitted_file(sub_feats.c_str(), &sub_tasks, &utt2feats_map, succ_num_one_time1);
				write_splitted_file(sub_targets.c_str(), &sub_tasks, &utt2targets_map, succ_num_one_time2);
				write_splitted_file(sub_utt2spk.c_str(), &sub_tasks, &utt2spk_map, succ_num_one_time3);
				if (succ_num_one_time1 == succ_num_one_time1 &&
					succ_num_one_time1 == succ_num_one_time3 && 
					succ_num_one_time2 == succ_num_one_time3){
					succ_num_one_time = succ_num_one_time2;
				}
				else{
					std::cout << "you should fix your data dir first .\n";
					exit(-1);

				}
				succ_num += succ_num_one_time;
				if (succ_num_one_time > 0){
					fprintf(stdout, "%s:%8d files have been written.\n", base_dir.c_str(), succ_num_one_time);
				}
				if (succ_num % 10000 == 0){
					fprintf(stdout, "%8d files have been written.\n",  succ_num);
				}
				data_protect.Leave();
			}
			neosmart::ResetEvent(h_event_start);
			neosmart::SetEvent(h_event_ack);
		}
		if (neosmart::CheckEvent(h_event_stop)){
			//do something
			neosmart::ResetEvent(h_event_stop);
			break;
		}

	}
	std::cout << "Total " << succ_num << " files have been written.\n";
	return 0;
}





string  ConvertWhite2SingleBlank(const char* original_string){
	string temp_string = original_string, final_string, next_token;
	if (temp_string.empty())
	{
		return temp_string;
	}
	istringstream temp_stream(temp_string);
	while (!temp_stream.eof())
	{
		temp_stream >> next_token;
		if (next_token != " " || next_token != "\t")
		{
			final_string = final_string + next_token;
			if (temp_stream.eof()){
				break;
			}
			else{
				final_string += " ";
			}
		}
	}

	return final_string;
}


string preprocess_string(const char* original_string){
	string temp_string = original_string, processed_string;

	if (temp_string.empty())
	{
		return temp_string;
	}
	processed_string = ConvertWhite2SingleBlank(temp_string.c_str());
	if (!processed_string.empty())
		return processed_string;
	return temp_string;

}
int FindSpecialPos(const char* original_str, const char* delimiter){
	int pos = -1;
	if (!original_str || !delimiter){
		return pos;
	}
	string deli_array = delimiter, test_str = original_str;
	vector<int> pos_array;
	int i = 0;
	while (i < deli_array.size())
	{
		string::size_type result_pos = test_str.find_first_of(deli_array[i]);
		if (result_pos != string::npos)
		{
			pos_array.push_back(result_pos);
		}
		i++;
	}
	std::vector<int>::iterator min_iter = min_element(pos_array.begin(), pos_array.end());
	if (min_iter != pos_array.end() && (*min_iter) > pos)
	{
		return (*min_iter);
	}
	return pos;
}

int split_line_special(const char* original_string, const char* delimiter, vector<string>* results, int split_times = 1){
	if (!original_string || !delimiter || !results){
		return -1;
	}
	string ori_str = original_string;
	string deli_str = delimiter, part_str = "";
	vector<string> delis;
	string::size_type sub_str_start = 0, sub_str_end = FindSpecialPos(ori_str.c_str(), deli_str.c_str());
	while (sub_str_end != string::npos)
	{
		results->push_back(ori_str.substr(sub_str_start, sub_str_end - sub_str_start));
		part_str = ori_str.substr(sub_str_end + 1);
		if (results->size() == split_times)
		{
			results->push_back(part_str);
			break;
		}
		sub_str_start = sub_str_end;
		sub_str_end = FindSpecialPos(part_str.c_str(), deli_str.c_str());

	}
	if (results->size()){
		return results->size();
	}
	return -1;
}



int UttToSpkInternal(const char* utt2spk_name, int min_data_entry_num){
	std::ifstream utt_is(utt2spk_name, ios::in);
	utt2spk_map.reserve(min_data_entry_num);
	utt2spk_array.reserve(min_data_entry_num);
	utt2spk_content.reserve(min_data_entry_num);
	int succ_num = 0;
	std::cout << "start to read " << utt2spk_name << std::endl;
	StdThread ps;
	try
	{
		std::string line;
		vector<string> line_units;
		StringManager string_manager;
		while (! utt_is.eof())
		{
			getline(utt_is, line, '\n');
			string_manager.set_string(line);
			if (string_manager.is_empty()){
				continue;
			}
			string_manager.SplitStringtoSpecializedSize(line.c_str(), " \t", &line_units, 1);
			if (line_units.size() == 2){
				string scp_key = line_units[0], scp_content = line_units[1];
				utt2spk_map.insert(std::pair<string, string>(scp_key, scp_content));
				utt2spk_keys_array.push_back(scp_key);
				succ_num++;
				if (succ_num % 10000 == 0){
					ps.Enter();
					std::cout << succ_num << " have been read.\n";
					ps.Leave();
				}
			}
			line_units.clear();

		}
	}
	catch (exception* e)
	{
		utt_is.close();
		std::cout << e->what() << std::endl;
	}
	ps.Enter();
	fprintf(stdout, "Total number file in %s:%8d\n", utt2spk_name, succ_num);
	ps.Leave();
	utt_is.close();
	
	return 0;
}



int GetNormalScp(const char* scp_name, FastMap* scp_map, StringArray* scp_keys, int min_data_entry_num){
	std::ifstream utt_is(scp_name, ios::in);
	if (scp_map){
		scp_map->reserve(min_data_entry_num);
	}
	if (scp_keys){
		scp_keys->reserve(min_data_entry_num);
	}
	StdThread print_cs;
	int succ_num = 0;
	try
	{
		std::string line;
		vector<string> line_units;
		while (!utt_is.eof())
		{
			getline(utt_is, line, '\n');
			line = preprocess_string(line.c_str());
			if (line.empty()) continue;
			split_line_special(line.c_str(), " \t", &line_units, 1);
			if (line_units.size() == 2){
				string scp_key = line_units[0], scp_content = line_units[1];
				bool is_good_entry = false;
				if (scp_map){
					scp_map->insert(std::pair<string, string>(scp_key, scp_content));
					is_good_entry = true;
				}
				if (scp_keys){
					scp_keys->push_back(line);
				}
				if (is_good_entry){
					succ_num++;
				}
				if (succ_num % 10000 == 0){
					print_cs.Enter();
					std::cout << succ_num << " files have been read.\n";
					print_cs.Leave();
				}
			}
			line_units.clear();

		}
	}
	catch (exception* e)
	{
		utt_is.close();
		std::cout << e->what() << std::endl;
	}
	utt_is.close();
	print_cs.Enter();
	std::cout << "total " << succ_num << " files have been read\n";
	print_cs.Leave();
	return 0;
}




int ReadandFixPdfs(const char* scp_name, FastMap* scp_map, StringArray* scp_keys, int min_data_entry_num){
	std::ifstream utt_is(scp_name, ios::in);
	if (scp_map){
		scp_map->reserve(min_data_entry_num);
	}
	if (scp_keys){
		scp_keys->reserve(min_data_entry_num);
	}
	StdThread print_cs;
	int succ_num = 0;
	try
	{
		std::string line;
		vector<string> line_units;
		while (!utt_is.eof())
		{
			getline(utt_is, line, '\n');
			line = preprocess_string(line.c_str());
			if (line.empty()) continue;
			split_line_special(line.c_str(), " \t", &line_units, 1);
			if (line_units.size() == 2){
				string scp_key = line_units[0], scp_content = line_units[1];
				bool is_good_entry = false;
				if (scp_map){
					scp_map->insert(std::pair<string, string>(scp_key, scp_content));
					is_good_entry = true;
				}
				if (scp_keys){
					scp_keys->push_back(line);
				}
				if (is_good_entry){
					succ_num++;
				}
				if (succ_num % 10000 == 0){
					print_cs.Enter();
					std::cout << succ_num << " files have been read.\n";
					print_cs.Leave();
				}
			}
			line_units.clear();

		}
	}
	catch (exception* e)
	{
		utt_is.close();
		std::cout << e->what() << std::endl;
	}
	utt_is.close();
	print_cs.Enter();
	std::cout << "total " << succ_num << " files have been read\n";
	print_cs.Leave();
	return 0;
}



int GetUtttoSpk(const char* utt2spk_name, int min_data_entry_num){
	if (!CheckFile(utt2spk_name) || min_data_entry_num < 0){
		std::cout << utt2spk_name << " " << "Needed file is not exist.\n";
		return -1;
	}
	return UttToSpkInternal(utt2spk_name, min_data_entry_num);
}



int SplitDataForBmuf(const char* utt2spk_name, 
	const char* feats_name, const char* targets_name, const char* data_dir, long long data_entry_num, int split_times){
	if (! CheckFile(utt2spk_name) || ! CheckFile(feats_name) || ! CheckFile(targets_name)){
		return -1;
	}
	
	std::thread read_utt2spk_task(GetUtttoSpk, utt2spk_name, data_entry_num);
	std::thread read_feats_task(GetNormalScp, feats_name, &utt2feats_map, nullptr, data_entry_num);
	std::thread read_targets_task(GetNormalScp, targets_name, &utt2targets_map, nullptr, data_entry_num);
	read_utt2spk_task.join();
	read_feats_task.join();
	read_targets_task.join();
	int utt_size = 0, feats_size = 0, targets_size = 0, sub_task_size, utt_array_size;
	utt_array_size = utt2spk_keys_array.size();
	utt_size = utt2spk_map.size();
	feats_size = utt2feats_map.size();
	targets_size = utt2targets_map.size();
	if (utt_size != feats_size || utt_size != targets_size || feats_size != targets_size || utt_size != utt_array_size){
		std::cout << "scp has different sizes, please check !\n";
		return -1;
	}
	else{
		sub_task_size = utt_size / split_times;
	}
	//start next 2 threads here
	producer_consumer_init();
	std::thread write_splitted_task(WriteScpFile);

	//start to store sub task to these 2 threads
	int sub_task_start, sub_task_end;
	
	sub_tasks.resize(sub_task_size);
	StdThread data_protect;
	task_info.Reset();
	task_info.sub_task_root_ = data_dir;
	task_info.sub_task_root_ += "//split" + boost::lexical_cast<string>(split_times);
	fprintf(stdout, "start to split data now.\n");
	for (int i = 0; i < split_times;i++)
	{
		sub_task_start = i*sub_task_size;
		sub_task_end = (i + 1)*sub_task_size ;
		if (sub_task_start > sub_task_end){
			std::cout << "ERR\n";
			exit(-1);
		}
		if (sub_task_end > utt2spk_keys_array.size()){
			sub_task_end = utt2spk_keys_array.size();
		}
		if (i == (split_times-1)){
			sub_task_end = utt2spk_keys_array.size();
			sub_tasks.resize(sub_task_end - sub_task_start);
		}
		std::copy(utt2spk_keys_array.begin() + sub_task_start, utt2spk_keys_array.begin() + sub_task_end, sub_tasks.begin());
		task_info.sub_task_index_ = i + 1;
		task_info.SetState(true);
		/////////////////// Here we should sent these data to other threads ////////////////////////
		neosmart::ResetEvent(h_event_ack);
		neosmart::SetEvent(h_event_start);
		//monitor the state of threads and send data to them
		neosmart::WaitForEvent(h_event_ack, -1);
		//////////////////////////////////////////////////////////////////////////
		sub_tasks.clear();
		sub_tasks.resize(sub_task_size);

	}
	//send ending information to the working thread and cancel all tasks
	neosmart::SetEvent(h_event_stop);
	write_splitted_task.join();
	//////////////////////////////////////////////////////////////////////////
	producer_consumer_stop();
	

	// when all data sent, all threads return and this function ends.
	return 0;

}


int main(int argc, char* argv[]){
	std::cout << "Start to split data for BMUF optimization\n";
	if (argc < 3){
		std::cout << argv[0] << " <data/train> <upper bound of original data, such as 10000000> <max splitted times, such as 12000>\n";
		return -1;
	}
	string data_root_dir = argv[1];
	long long int data_entries_upbound = std::atoi(argv[2]);
	int split_times = std::atoi(argv[3]);
	string utt2spk_file_name, feats_file_name, targets_file_name;
	if (! data_root_dir.empty()){
		utt2spk_file_name = data_root_dir + "\/utt2spk";
		feats_file_name = data_root_dir + "\/feats.scp";
		targets_file_name = data_root_dir + "\/targets.scp";

	}
	
	try
	{
		SplitDataForBmuf(utt2spk_file_name.c_str(), feats_file_name.c_str(), targets_file_name.c_str(), data_root_dir.c_str(), 
			data_entries_upbound, split_times);
		
	}
	catch (exception* e)
	{
		std::cout << e->what() << std::endl;
		std::cout << " failed\n";
	}
	return 0;
}
